Note: CodeRabbit is enabled on this repository as a convenience for maintainers
and contributors. Use your best judgment when considering its review comments and
suggestions — a suggested change may be inadequate, unnecessary, or safe to ignore.
Contributors are not expected to address every comment. Human reviews are what
ultimately matter for merging.

Summary

This PR implements cudax::coop::reduce for arbitrary warp-groupings within a CUDA block (fixes #9160). It adds a warp-group reduction path that performs a two-stage reduction across warps inside a block, adjusts the public wrapper to take the reduction functor by value, and adds Catch2 tests validating multi-warp reductions. Several supporting group/mapping/synchronizer fixes were included to ensure correct rank/mapping and barrier initialization when units are not at thread level.

Core Implementation Changes

• cudax/include/cuda/experimental/__coop/reduce.cuh
  • Added an overload of __reduce_impl constrained to groups with unit_type == warp_level and level_type == block_level.
  • Two-stage reduction:
    • Per-warp reductions using cub::WarpReduce with per-warp TempStorage; each warp writes one partial into shared memory.
    • Synchronize, then reduce per-warp partials using a root-warp cub::WarpReduce; only the group root returns a cuda::std::optional containing the final reduced value, other threads return nullopt.
  • Public wrapper signature changed from reduce(_Group, _Tp (&)[_Np], _RedFn&&) → reduce(_Group, _Tp (&)[_Np], _RedFn) (functor passed by value and forwarded into the impl).

Supporting Changes

• cudax/include/cuda/experimental/__group/queries.cuh

  • In the __rank_query_group non-matching branch, the linear rank computation now multiplies the group-unit rank by the unit count derived from __group.hierarchy() and adds the unit-local rank (uses __count_query/__rank_query from the hierarchy).

• cudax/include/cuda/experimental/__group/group.cuh

  • Lane-mask validation in __do_mapping made conditional on _Unit: when _Unit == thread_level assert lane inclusion and popcount constraints; otherwise require lane mask == all lanes (previously always validated per-lane).

• cudax/include/cuda/experimental/__group/synchronizer/barrier_synchronizer.cuh

  • make_instance now precomputes per-Unit thread metadata (__nthread_in_unit, __thread_rank_in_unit) and derives them from the parent hierarchy when _Unit != thread_level.
  • Barrier initialization guard tightened to require mapping valid, mapping rank == 0, and __thread_rank_in_unit == 0.
  • Barrier init count scaled to __mapping_result.count() * __nthread_in_unit to account for multiple threads per unit.

Tests

• cudax/test/coop/reduce/warps_within_block.cu (new)

  • Adds ReduceKernel and parametrized tests that create groups of nwarps_in_group warps and verify cudax::coop::reduce across multiple warps.
  • Tests:
    • Integral types with integer reduction operators (plus/mul/bitwise/min/max) — exact equality.
    • Floating-point types (float/double) with FP ops (plus/mul/min/max) — relative tolerance.
    • Per-thread input sizes 1–4.
  • Kernel runtime config supports block size = (nwarps_in_group + 2) * warp_size; only group-member threads participate in the reduction and only group root writes output.

• cudax/test/CMakeLists.txt

  • Registers new test target coop.reduce.warps_within_block.

API / ABI Surface Changes

• Added internal overload: __reduce_impl(_Group, _Tp (&)[_Np], _RedFn) for warp-group reductions.
• Changed public reduce signature to take the reduction functor by value (_RedFn), instead of an rvalue reference.

Review Notes / Areas to Verify

• Shared-memory layout and sizing for per-warp partials and correctness of root-warp selection (ensure alignment and no bank-conflict issues).
• Correctness and performance across varying warp-group sizes and SM architectures (two-stage approach and root-warp reduction).
• Impact of changing public reduce functor parameter to pass-by-value on forwarding/move semantics.
• Barrier initialization and mapping correctness when units contain multiple threads (ensure counts/ranks align with mapping and that lane-mask validation changes are appropriate).

important:

Walkthrough

Adds a warp-scoped two-stage cooperative reduce overload, conditions group mapping/queries on unit hierarchy and adjusts barrier init for per-unit threads, and adds parametrized CUDA tests validating multi-warp reductions across types, ops, and item counts.

Changes

Cooperative reduce for warps within block

Assessment against linked issues

Possibly related PRs

• NVIDIA/cccl#9167: Adds other __reduce_impl overloads and related tests; closely aligned with multi-overload reduce work.
• NVIDIA/cccl#9154: Earlier changes to __reduce_impl overloads and reduce wrapper plumbing; directly related.
• NVIDIA/cccl#9264: Similar conditionalization of lane-mask validation in group.cuh; related to group mapping changes.

Suggested labels

cudax

Suggested reviewers

• andralex
• miscco
• pciolkosz

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

🥳 CI Workflow Results

🟩 Finished in 35m 02s: Pass: 100%/55 | Total: 8h 09m | Max: 35m 00s | Hits: 71%/46192

See results here.